#!/usr/bin/env python3 
from math import sqrt
import torch

class ChunkGraphGCNNormalizer:
    def __init__(self, nvertex:int ):
        self.nvertex = nvertex
    
    def __call__(self, edge_list: torch.Tensor, weight_list: torch.Tensor):
        
        assert edge_list.size(0) == 2

        norm_ew = torch.ops.helpers.gcn_norm_coo(edge_list.cuda(), weight_list.cuda(), self.nvertex)
        norm_rw = torch.ops.helpers.gcn_norm_coo(edge_list.flip(0).contiguous().cuda(), weight_list.cuda(), self.nvertex)

        if (norm_ew == norm_rw).all():
            return norm_ew.cpu()
        else:
            return norm_ew.cpu(), norm_rw.cpu()
            

